##########################################
## A Function to Combine Estimates from MI
##########################################

# Borrowed from Carlisle Rainey's public GitHub
# source: https://gist.github.com/carlislerainey/9963264

comb.mi <- function(models) {
  
  # Arguments
  #   models:  a list of models, one estimated on each of m MI data sets.
  
  # Note: I borrow the notation below mostly from Rubin's *Multiple Imputation 
  # for Nonresponse in Surveys*.
  
  # count the number of imputations
  n.imp <- length(models)
  
  # calculate the average of the coefficients (Eqn 3.1.2, p. 76)
  mi.est <- NULL
  for (i in 1:n.imp) {
    mi.est <- rbind(mi.est, coef(models[[i]]))
  }
  Q.bar <- apply(mi.est, 2, mean)
  # store estimates
  comb.est <- Q.bar
  
  # calculate the average of the covariances (Eqn 3.1.3, p. 76)
  U.bar <- NULL
  U.bar <- vcov(models[[1]])
  for (i in 2:n.imp) {
    U.bar <- U.bar + vcov(models[[i]])
  }
  U.bar <- U.bar/n.imp
  
  # calculate the variance *between* the MI data sets (Eqn 3.1.4, p. 76)
  B <- NULL
  B <- (coef(models[[1]]) - Q.bar)%*%t(coef(models[[1]]) - Q.bar)
  for (i in 2:n.imp) {
    B <- B + (coef(models[[i]]) - Q.bar)%*%t(coef(models[[i]]) - Q.bar)
  }
  B <- B/(n.imp - 1)
  
  # calculate the *total* covariance (Eqn 3.1.5, p. 76)
  comb.cov <- U.bar + (1 + 1/(n.imp))*B # T iin Rubin's notation
  
  # calculate the standard errors
  comb.se <- sqrt(diag(comb.cov))
  
  # add the relevant quantities to a list and return it
  res <- list(comb.est = comb.est, comb.cov = comb.cov, comb.se = comb.se)
  return(res)
  
  # print the estimates and standard errors
  print.it <- cbind(comb.est, comb.se)
  rownames(print.it) <- names(comb.est)
}

